import tkinter as tk
from tkinter import ttk
import random
import math
import numpy as np              # 添加NumPy用于信号处理
from tkinter import filedialog  # 文件选择对话框
from ui.BasePage import BasePage

class DataPreprocessingPage(BasePage):
    """数据预处理页面（整合版）"""

    def __init__(self, parent, controller):
        super().__init__(parent, controller)
        
        self.local_data = None  # 存储加载的本地数据
        self.load_data()        # 加载数据

        """ 创建外层水平方向的 PanedWindow，实现左右分割 """
        outer_paned = ttk.PanedWindow(self.frame, orient=tk.HORIZONTAL)
        outer_paned.pack(fill=tk.BOTH, expand=True)

        """ 创建左侧的 Frame 用于放置预处理选项 """
        left_frame = ttk.Frame(outer_paned, padding=(10, 10), width=250)
        left_frame.pack_propagate(False)  # 禁止子组件影响父容器尺寸
        outer_paned.add(left_frame)

        # 数据归一化部分
        data_normalize_frame = ttk.LabelFrame(
            left_frame,
            text="数据归一化",
        )
        data_normalize_frame.pack(fill=tk.X, pady=5)
        
        norm_choices = [
            "不归一化",
            "Min-MAX归一化",
            "标准化（Z-score）",
            "Savitzky-Golay平滑滤波",
        ]
        self.norm_var = tk.StringVar()
        self.norm_var.set(norm_choices[1])
        for choice in norm_choices:
            radio = ttk.Radiobutton(
                data_normalize_frame, text=choice, variable=self.norm_var, value=choice
            )
            radio.pack(anchor=tk.W, pady=2)

        # 水平分隔线
        sep1 = ttk.Separator(left_frame, orient=tk.HORIZONTAL)
        sep1.pack(fill=tk.X, padx=5, pady=5)

        # 特征提取部分
        feature_extract_frame = ttk.LabelFrame(
            left_frame,
            text="特征提取",
        )
        feature_extract_frame.pack(fill=tk.X, pady=5)
        
        feat_choices = ["原始信号", "FFT频谱", "包络谱", "小波变换"]
        self.feat_var = tk.StringVar()
        self.feat_var.set(feat_choices[0])
        for choice in feat_choices:
            radio = ttk.Radiobutton(
                feature_extract_frame, text=choice, variable=self.feat_var, value=choice
            )
            radio.pack(anchor=tk.W, pady=2)

        # 水平分隔线
        sep2 = ttk.Separator(left_frame, orient=tk.HORIZONTAL)
        sep2.pack(fill=tk.X, padx=5, pady=5)

        # 训练集/测试集划分部分
        split_frame = ttk.LabelFrame(
            left_frame,
            text="训练集/测试集划分",
        )
        split_frame.pack(fill=tk.X, pady=5)

        split_options_frame = ttk.Frame(split_frame)
        split_options_frame.pack(fill=tk.X, pady=5)

        self.split_scale = ttk.Scale(
            split_options_frame, from_=0.01, to=0.99, orient=tk.HORIZONTAL, length=150
        )
        self.split_scale.set(0.2)
        self.split_scale.pack(side=tk.LEFT, padx=5)

        self.split_percent_var = tk.StringVar(value="20%")
        split_percent_label = ttk.Label(
            split_options_frame, textvariable=self.split_percent_var
        )
        split_percent_label.pack(side=tk.LEFT, padx=5)

        self.split_scale.bind("<Motion>", self.update_percentage_label)
        self.split_scale.bind("<ButtonRelease-1>", self.update_percentage_label)

        # 应用预处理按钮
        self.preprocess_button = ttk.Button(
            split_frame, text="应用预处理", style="Info.TButton"
        )
        self.preprocess_button.pack(anchor=tk.CENTER, pady=5)
        self.preprocess_button.bind("<Button-1>", self.apply_preprocessing)

        """ 创建右侧垂直方向的 PanedWindow """
        right_paned = ttk.PanedWindow(
            outer_paned, orient=tk.VERTICAL
        )
        outer_paned.add(right_paned)

        # 顶部控制区域
        control_frame = ttk.LabelFrame(
            right_paned,
            text="控制区域",
        )
        control_frame.pack_propagate(False)  # 禁止内容拉伸高度
        control_frame.configure(height=120)  # 固定高度
        right_paned.add(control_frame)

        # 更新可视化按钮
        self.update_button = ttk.Button(
            control_frame,
            text="更新可视化",
            style="Info.TButton"
        )
        self.update_button.pack(pady=5)
        self.update_button.bind("<Button-1>", self.update_visualization)

        # 可视化类型选择区域
        visualize_type_frame = ttk.Frame(control_frame)
        visualize_type_frame.pack(fill=tk.X, pady=5)

        visualize_type_label = ttk.Label(
            visualize_type_frame, text="可视化类型："
        )
        visualize_type_label.pack(side=tk.LEFT, padx=5)

        self.visualize_type_combobox = ttk.Combobox(
            visualize_type_frame,
            values=["原始信号", "FFT频谱", "包络谱", "时频图", "特征分布"],
            width=15,
            state="readonly",
        )
        self.visualize_type_combobox.set("特征分布")
        self.visualize_type_combobox.pack(side=tk.LEFT, padx=5)
        self.visualize_type_combobox.bind("<<ComboboxSelected>>", self.update_visualization)

        fault_type_label = ttk.Label(
            visualize_type_frame, text="故障类型："
        )
        fault_type_label.pack(side=tk.LEFT, padx=5)

        self.fault_type_combobox = ttk.Combobox(
            visualize_type_frame,
            values=["正常", "内圈故障", "外圈故障", "滚动体故障", "保持架故障"],
            width=10,
            state="readonly",
        )
        self.fault_type_combobox.set("正常")
        self.fault_type_combobox.pack(side=tk.LEFT, padx=5)
        self.fault_type_combobox.bind("<<ComboboxSelected>>", self.update_visualization)

        # 底部可视化显示区域
        display_frame = ttk.LabelFrame(
            right_paned,
            text="可视化显示",
        )
        right_paned.add(display_frame)

        # Canvas自适应窗口大小
        self.display_canvas = tk.Canvas(display_frame, bg="white")
        self.display_canvas.pack(fill=tk.BOTH, expand=True)

        # 初始化数据
        self.sample_rate = 2048  # 采样率Hz
        self.signal_length = 2048  # 信号长度
        self.time_data = np.linspace(0, 1, self.signal_length)
        
        self.original_data = self.generate_signal()
        self.processed_data = self.original_data.copy()
        
        # 绑定窗口大小变化事件
        self.frame.bind("<Configure>", self.on_frame_configure)
        
        # self.update_visualization()
        
    def load_data(self, event=None):
        """ 加载本地数据文件 用于加载Numpy数据 """
        file_path = filedialog.askopenfilename(
            title="选择数据文件",
            filetypes=[("NumPy文件", "*.npy"), ("MAT文件", "*.mat"), ("所有文件", "*.*")]
        )
        
        if not file_path:
            return
            
        try:
            # 根据文件扩展名选择加载方法
            if file_path.endswith('.npy'):
                self.local_data = np.load(file_path)
            elif file_path.endswith('.mat'):
                # 假设这是MATLAB格式的文件，需要scipy.io.loadmat读取
                from scipy.io import loadmat
                mat_data = loadmat(file_path)
                # 查找数据数组（假设第一个非元数据的数组是我们需要的）
                for key in mat_data.keys():
                    if not key.startswith('__'):
                        self.local_data = mat_data[key].flatten()
                        break
            else:
                # 尝试作为文本文件读取
                self.local_data = np.loadtxt(file_path)
                
            # 确保数据是一维数组
            if self.local_data.ndim > 1:
                self.local_data = self.local_data.flatten()
                
            # 更新信号长度以匹配加载的数据
            self.signal_length = len(self.local_data)
            self.time_data = np.linspace(0, self.signal_length/self.sample_rate, self.signal_length)
            
            # 更新原始数据
            self.original_data = self.generate_signal()
            self.processed_data = self.original_data.copy()
            
            # 更新可视化
            self.update_visualization()
            
            print(f"成功加载数据，长度: {len(self.local_data)}")
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            
    def update_percentage_label(self, event):
        """更新训练集划分百分比显示"""
        value = self.split_scale.get()
        self.split_percent_var.set(f"{int(value*100)}%")

    def generate_signal(self):
        """生成信号（优先使用加载的本地数据）"""
        # 如果已加载本地数据，则使用本地数据
        if hasattr(self, 'local_data') and self.local_data is not None:
            # 调整本地数据长度以匹配信号长度
            if len(self.local_data) > self.signal_length:
                return self.local_data[:self.signal_length]
            else:
                # 如果数据不足，用零填充或循环数据
                padded_data = np.zeros(self.signal_length)
                padded_data[:len(self.local_data)] = self.local_data
                return padded_data
                
        # 否则生成随机信号（默认情况）
        fault_type = self.fault_type_combobox.get()
        t = self.time_data
        
        # 基础信号（更大幅度）
        signal = 2 * np.sin(2 * np.pi * 50 * t)  # 基频（振幅2）
        signal += 1 * np.sin(2 * np.pi * 100 * t)  # 二次谐波（振幅1）
        signal += 0.5 * np.sin(2 * np.pi * 150 * t)  # 三次谐波（振幅0.5）
        noise = np.random.normal(0, 0.3, self.signal_length)  # 更大噪声
        signal += noise
        
        # 故障特征（增强幅度）
        if fault_type == "内圈故障":
            signal += 1.5 * np.sin(2 * np.pi * 100 * t) * np.exp(-t*2)  # 更强冲击
        elif fault_type == "外圈故障":
            signal *= (1 + 0.6 * np.sin(2 * np.pi * 10 * t))  # 更强调制
            signal += 1.2 * np.sin(2 * np.pi * 80 * t)  # 更高幅值特征
        elif fault_type == "滚动体故障":
            impacts = np.zeros(self.signal_length)
            for i in range(8):  # 更多冲击
                pos = random.randint(0, self.signal_length-80)
                impact = np.exp(-np.linspace(0, 3, 80)) * np.sin(2 * np.pi * 60 * np.linspace(0, 0.1, 80)) * 2  # 更高冲击幅度
                impacts[pos:pos+80] += impact
            signal += impacts
        elif fault_type == "保持架故障":
            signal += 1.0 * np.sin(2 * np.pi * 20 * t)  # 更高幅值特征
        
        return signal

    def apply_preprocessing(self, event=None):
        """应用预处理并更新数据"""
        norm_method = self.norm_var.get()
        feat_method = self.feat_var.get()
        
        # 复制原始数据
        self.processed_data = self.original_data.copy()
        
        # 应用归一化
        if norm_method == "Min-MAX归一化":
            self.processed_data = (self.processed_data - np.min(self.processed_data)) / (np.max(self.processed_data) - np.min(self.processed_data)) * 2 - 1
        elif norm_method == "标准化（Z-score）":
            self.processed_data = (self.processed_data - np.mean(self.processed_data)) / np.std(self.processed_data)
        elif norm_method == "Savitzky-Golay平滑滤波":
            # 简化的Savitzky-Golay滤波
            window_size = 11
            order = 3
            half = window_size // 2
            b = np.polyfit(np.arange(window_size), np.ones(window_size), order)
            for i in range(half, len(self.processed_data)-half):
                self.processed_data[i] = np.sum(self.processed_data[i-half:i+half+1] * b)
        
        # 应用特征提取
        if feat_method == "FFT频谱":
            fft_data = np.fft.rfft(self.processed_data)
            self.processed_data = np.abs(fft_data) / (self.signal_length/2)
            self.time_data = np.fft.rfftfreq(self.signal_length, 1/self.sample_rate)
        elif feat_method == "包络谱":
            from scipy.signal import hilbert
            analytic = hilbert(self.processed_data)
            envelope = np.abs(analytic)
            fft_envelope = np.fft.rfft(envelope)
            self.processed_data = np.abs(fft_envelope) / (self.signal_length/2)
            self.time_data = np.fft.rfftfreq(self.signal_length, 1/self.sample_rate)
        elif feat_method == "小波变换":
            # 简化小波变换（降采样）
            self.processed_data = self.processed_data[::2]
            self.time_data = self.time_data[::2]
        
        # 更新可视化
        self.update_visualization()

    def update_visualization(self, event=None):
        """更新可视化显示"""
        visualize_type = self.visualize_type_combobox.get()
        data_to_plot = self.processed_data
        x_label = "时间 (秒)"
        y_label = "幅值"
        
        if visualize_type == "原始信号":
            data_to_plot = self.original_data
            y_label = "幅值"
        elif visualize_type == "FFT频谱":
            x_label = "频率 (Hz)"
            y_label = "幅值"
        elif visualize_type == "包络谱":
            x_label = "频率 (Hz)"
            y_label = "幅值"
        elif visualize_type == "时频图":
            self.draw_time_frequency()
            return
        elif visualize_type == "特征分布":
            self.draw_feature_distribution()
            return
        
        self.draw_signal(data_to_plot, x_label, y_label)

    def draw_signal(self, data, x_label, y_label):
        """绘制信号波形（优化坐标映射）"""
        self.display_canvas.delete("all")
        canvas_width = self.display_canvas.winfo_width()
        canvas_height = self.display_canvas.winfo_height()
        margin = max(30, canvas_width // 20)
        
        # 计算数据范围
        data_min = np.min(data)
        data_max = np.max(data)
        data_range = data_max - data_min + 1e-10  # 避免除以零
        
        # 绘制坐标轴
        self.display_canvas.create_line(margin, margin, margin, canvas_height-margin, width=2)  # Y轴
        self.display_canvas.create_line(margin, canvas_height-margin, canvas_width-margin, canvas_height-margin, width=2)  # X轴
        
        # 绘制波形（使用线性映射充满画布高度）
        points = []
        data_len = len(data)
        for i in range(data_len):
            x = margin + i * (canvas_width - 2*margin) / (data_len - 1)  # X坐标均匀分布
            y = canvas_height - margin - (data[i] - data_min) * (canvas_height - 2*margin) / data_range  # Y坐标充满高度
            points.append(x)
            points.append(y)
        
        if points:
            self.display_canvas.create_line(points, fill="blue", width=2)
        
        # 添加标签和标题
        self.display_canvas.create_text(canvas_width-margin, canvas_height-margin+20, text=x_label, font=("宋体", 10))
        self.display_canvas.create_text(margin-20, (margin+canvas_height-margin)//2, text=y_label, font=("宋体", 10), angle=90)
        title = f"{self.fault_type_combobox.get()} - {self.visualize_type_combobox.get()}"
        self.display_canvas.create_text(canvas_width//2, margin//2, text=title, font=("宋体", 12, "bold"))
        
    def draw_time_frequency(self):
        """绘制时频图"""
        self.display_canvas.delete("all")
        canvas_width = self.display_canvas.winfo_width()
        canvas_height = self.display_canvas.winfo_height()
        margin = max(30, canvas_width // 20)
        
        # 绘制坐标轴
        self.display_canvas.create_line(margin, margin, margin, canvas_height-margin, width=2)
        self.display_canvas.create_line(margin, canvas_height-margin, canvas_width-margin, canvas_height-margin, width=2)
        
        # 绘制时频图（简化为颜色块）
        time_bins = 20
        freq_bins = 15
        for i in range(time_bins):
            for j in range(freq_bins):
                energy = random.random() * (0.5 + 0.5*abs(math.sin(i*freq_bins + j)))
                r = int(255 * energy)
                g = int(128 * energy)
                b = int(64 * energy)
                color = f"#{r:02x}{g:02x}{b:02x}"
                
                x1 = margin + i * (canvas_width - 2*margin) / time_bins
                y1 = canvas_height - margin - (j+1) * (canvas_height - 2*margin) / freq_bins
                x2 = x1 + (canvas_width - 2*margin) / time_bins
                y2 = y1 + (canvas_height - 2*margin) / freq_bins
                
                self.display_canvas.create_rectangle(x1, y1, x2, y2, fill=color, outline="")
        
        # 添加标签
        self.display_canvas.create_text(canvas_width-margin, canvas_height-margin+20, text="时间 (秒)", font=("宋体", 10))
        self.display_canvas.create_text(margin-20, (margin+canvas_height-margin)//2, text="频率 (Hz)", font=("宋体", 10), angle=90)
        title = f"{self.fault_type_combobox.get()} - 时频图"
        self.display_canvas.create_text(canvas_width//2, margin//2, text=title, font=("宋体", 12, "bold"))

    def draw_feature_distribution(self):
        """绘制特征分布图"""
        self.display_canvas.delete("all")
        canvas_width = self.display_canvas.winfo_width()
        canvas_height = self.display_canvas.winfo_height()
        margin = max(30, canvas_width // 20)
        
        # 绘制坐标轴
        self.display_canvas.create_line(margin, margin, margin, canvas_height-margin, width=2)
        self.display_canvas.create_line(margin, canvas_height-margin, canvas_width-margin, canvas_height-margin, width=2)
        
        # 生成特征值
        features = ["均值", "方差", "峰值", "峭度", "偏度"]
        values = [random.uniform(0, 1) for _ in range(5)]
        
        # 绘制柱状图
        bar_width = (canvas_width - 2*margin) / (len(features) * 1.2)
        for i, val in enumerate(values):
            x = margin + i * bar_width * 1.2 + bar_width/2
            y = canvas_height - margin - val * (canvas_height - 2*margin)
            self.display_canvas.create_rectangle(x-bar_width/2, y, x+bar_width/2, canvas_height-margin, fill="blue", outline="")
            self.display_canvas.create_text(x, canvas_height-margin+15, text=features[i], font=("宋体", 9))
            self.display_canvas.create_text(x, y-10, text=f"{val:.2f}", font=("宋体", 9))
        
        # 添加标签
        self.display_canvas.create_text(canvas_width-margin, canvas_height-margin+20, text="特征", font=("宋体", 10))
        self.display_canvas.create_text(margin-20, (margin+canvas_height-margin)//2, text="值", font=("宋体", 10), angle=90)
        title = f"{self.fault_type_combobox.get()} - 特征分布"
        self.display_canvas.create_text(canvas_width//2, margin//2, text=title, font=("宋体", 12, "bold"))

    def on_frame_configure(self, event):
        """窗口大小变化时更新可视化"""
        # self.update_visualization()